from scapy.all import *
import gmssl.sm3 as sm3
from gmssl.sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT
import random
import time
from time_master_send import chk_sum
from scapy.layers.inet import IP, TCP
from scapy.utils import wrpcap
from scapy.sendrecv import sniff
import operator

# 字符串转二进制
def str2bin(s):  
    return ''.join([bin(ord(c)).replace('0b', '') for c in s])

def xor_block(left, right):
    return map(operator.xor, left, right)

def to_bytes(data):
    if isinstance(data, str):
        if len(data) % 2 != 0:
            data += '0'
            print("add '0' at end,amended: ", end="")
            print(data)
        return bytes().fromhex(data)
    elif isinstance(data, list):
        return bytes(data)
    else:
        print("only 'str' or 'list' is valid!")
        return None


def to_hex(data):
    if isinstance(data, list):
        bytes_data = to_bytes(data)
        return bytes_data.hex()
    elif isinstance(data, bytes):
        return data.hex()
    else:
        print("only 'list' or 'bytes' is valid!")
        return None


def encrypt(massage, key):
    padding = "1111111111111111"
    key += padding
    key = bytes(key, encoding='utf-8')
    value = massage.encode(encoding='utf-8')  # bytes类型
    sm4 = CryptSM4()
    sm4.set_key(key, SM4_ENCRYPT)
    encrypt_value = sm4.crypt_ecb(value)  # bytes类型
    return to_hex(encrypt_value)  # 返回str


def decrypt(massage, key):
    key = bytes(key, encoding='utf-8')
    sm4 = CryptSM4()
    sm4.set_key(key, SM4_DECRYPT)
    decrypt_value = sm4.crypt_ecb(to_bytes(massage))  # bytes类型
    return decrypt_value.decode(encoding='utf-8')  # 返回str


# 字符串 data 进行hash 得到定长字符串
def my_hash(data):
    b = str(data)
    if len(b) % 2 != 0:
        b += "0"
    ans_list = []
    for i in range(int(len(b) / 2)):
        temp = b[i * 2:i * 2 + 2].encode('utf-8')
        ans_list.append(sm3.sm3_kdf(temp, 2))
    ans = ""
    for i in ans_list:
        ans += i
    end = ans[len(ans) - 1] + ans[0] + ans[int(len(ans) / 3)] + ans[int(len(ans) / 2)]
    return end


# 时间戳转换函数
def timestamp2time(timestamp):
    tmp = time.localtime(timestamp)  # time.localtime()格式化时间戳为本地时间
    my_time = time.strftime("%Y-%m-%d %H:%M:%S", tmp)  # 将本地时间格式化为字符串
    return my_time


# 抓包，获取包信息，存包
def get_print(pkt):
    with open("./temp.txt", "w", encoding='utf-8') as f:
        while True:
            try:
                if "?" in pkt.sprintf("{Raw:%Raw.load%}").split(r"\r\n")[0]:
                    wrpcap("sor.cap", pkt, append=True)
                    f.write(str(pkt['Ethernet'].dst) + '^')
                    f.write(str(pkt['Ethernet'].src) + '^')
                    f.write(str(pkt['Ethernet'].type) + '^')
                    f.write(str(pkt['IP'].version) + '^')
                    f.write(str(pkt['IP'].ihl) + '^')
                    f.write(str(pkt['IP'].tos) + '^')
                    f.write(str(pkt['IP'].len) + '^')
                    f.write(str(pkt['IP'].id) + '^')
                    f.write(str(pkt['IP'].flags) + '^')
                    f.write(str(pkt['IP'].frag) + '^')
                    f.write(str(pkt['IP'].ttl) + '^')
                    f.write(str(pkt['IP'].proto) + '^')
                    f.write(str(pkt['IP'].chksum) + '^')
                    f.write(str(pkt['IP'].src) + '^')
                    f.write(str(pkt['IP'].dst) + '^')
                    f.write(str(pkt['TCP'].sport) + '^')
                    f.write(str(pkt['TCP'].dport) + '^')
                    f.write(str(pkt['TCP'].seq) + '^')
                    f.write(str(pkt['TCP'].ack) + '^')
                    f.write(str(pkt['TCP'].dataofs) + '^')
                    f.write(str(pkt['TCP'].reserved) + '^')
                    f.write(str(pkt['TCP'].flags) + '^')
                    f.write(str(pkt['TCP'].window) + '^')
                    f.write(str(pkt['TCP'].chksum) + '^')
                    f.write(str(pkt['TCP'].urgptr) + '^')
                    f.write(str(pkt['TCP'].options) + '^')
                    f.write(str(pkt['Raw'].load))
                    break
                else:
                    print("抓包失败")
            except():
                exit(-1)

# secret_len参数是加密后的16进制位数
def get_packet_nub(secret_len):  # 返回需要的包的数量,和字数序列
    k = 0
    group = []
    while secret_len > 0:
        tf = random.randint(15, 20)
        if tf > secret_len:
            group.append(secret_len)
        else:
            group.append(tf)
        secret_len = secret_len - tf
        k = k + 1
    return k, group


def secret_split(secret):  # 数据分片
    ki = len(secret)
    su, group = get_packet_nub(ki)
    en = 0
    g = []
    for i in group:
        beg = en
        en = en + i
        g.append(secret[beg:en])
    return g, su, group

def keyxor(i,temp):
    hex_string = bytes(xor_block(bytes(i,"utf-8"), bytes(temp,"utf-8"))).hex()
    temp_list = ['a', 'b', 'c', 'd', 'e', 'f']
    temp = []
    for i in range(len(hex_string)):
        if hex_string[i] in temp_list:
            temp.append(temp_list.index(hex_string[i]))
        else:
            temp.append(hex_string[i])
    end = str(temp[len(hex_string) - 1]) + str(temp[0]) + str(temp[int(len(hex_string) / 3)]) + str(
        temp[int(len(hex_string) / 2)])
    return end

def get_pool_packets(num):
    packets = []
    pool = rdpcap("pool.cap")
    for pkt in pool:
        if "?" in pkt.sprintf("{Raw:%Raw.load%}").split(r"\r\n")[0]:
            packets.append(pkt)

    chk = []
    for pkt in packets:
        chk.append(str(pkt['TCP'].chksum))
    chk = list(set(chk))

    p = []
    for pkt in packets:
        if str(pkt["TCP"].chksum) in chk:
            p.append(pkt)
            chk.remove(str(pkt["TCP"].chksum))
            
    packets = []
    index = random.sample(range(1, len(p)), num)
    for i in index:
        packets.append(p[i])

    return packets


def secret_insert(data, seq, port, src_ip, dst_ip, end, key):
    with open("./temp.txt", "r", encoding="utf-8") as f:
        text = f.read()
        tf = text.split('^')

    tf_temp = tf[-1].split("?")
    data = tf_temp[0][2:] + "?c=un&m=" + data + "&" + tf_temp[1][:-1]
    data_list = data.split(r'\r\n')
    temp = ''
    for i in data_list:
        if i != '':
            temp += i
            temp += '\r\n'
    temp += '\r\n'
    data = temp
    if end is True:
        pkt = (IP(id=int(tf[7]), src=src_ip, dst=dst_ip)
               / TCP(sport=port, dport=80, seq=seq, ack=int(keyxor(chk_sum(port),key)))
               / data)
    else:
        pkt = (IP(id=int(tf[7]), src=src_ip, dst=dst_ip)
               / TCP(sport=port, dport=80, seq=seq, ack=int(keyxor(chk_sum(seq),key)))
               / data)
    # TCP数据包的seq字段值不能重复
    # send(pkt, iface="VMware Virtual Ethernet Adapter for VMnet8")
    wrpcap("dst.cap", pkt, append=True)
    return pkt


def make_packets(message: str = "This is a Secret message.", key: str = "password",
                 src: str = "192.168.31.155", dst: str = "192.168.31.155",id:int=852353):
    # 构造秘密信息
    secret = encrypt(message, key)  # 返回str
    #校验信息
    hash_value = my_hash(secret)
    #秘密信息分片
    secret_group, packet_num, group = secret_split(secret)
    num=int((packet_num-1)*3/8)
    if num>len(secret_group[0]):
        num=len(secret_group[0])
    time_data=secret_group[0][:num]
    temp=""
    for i in range(num):
        temp+="0"
    secret_group[0]=temp+secret_group[0][num:]
    # 解密部分
    # de_msg = decrypt(secret, key)
    # print('解密结果为:\n', de_msg)

    # 嵌入秘密信息
    #取出待嵌入信息的数据包
    packets = get_pool_packets(packet_num + 1)
    #数据包
    seq = id
    port = 12315
    pkt_list = []

    for i in range(packet_num):
        seq += 20
        port += 1
        get_print(packets[i])
        temp = secret_insert(secret_group[i], seq, port, src, dst,False, key)
        pkt_list.append(temp)
    #尾包
    seq = id
    port = 12310
    get_print(packets[packet_num])
    pkt_end = secret_insert(hash_value, seq, port, src, dst, True, key)
    #返回造好的包
    return packet_num, pkt_list, pkt_end,time_data

'''
def str_list2bytes(data):
    if isinstance(data, str):
        if len(data) % 2 != 0:
            data += '0'
            print("add '0' at end,amended: ", end="")
            print(data)
        return bytes().fromhex(data)
    elif isinstance(data, list):
        return bytes(data)
    else:
        print("only 'str' or 'list' is valid!")
        return None
    
def get_packet(interface, port, counts):
    sniff(iface=interface, prn=get_print, lfilter=lambda p: "GET" in str(p), filter=port, count=counts)

def make_packets(message: str = "This is a Secret message.", key: str = "p",
                 src: str = "192.168.188.66", dst: str = "192.168.188.90",
                 hash_value: str = "11111111111111111111"):
    # interface = ' Intel(R) Wireless-AC 9461 #2'  # 我的抓包网卡,
    # interface = "Intel(R) Wireless-AC 9560 160MHz"
    # port = "tcp port 80"
    # counts = 1
    # get_packet(interface, port, counts)  # 抓包

    # 构造秘密信息
    # tl = str2bin(key)
    secret = encrypt(message, key)  # 返回str
    secret_group, packet_num, group = secret_split(secret)

    # print("秘钥二进制：\n", tl)
    # print('加密结果为:\n', len(secret), secret)
    # print('需要包的数量：\n', packet_num)
    # print('每一个包数据数量：\n', group)
    # print('存入的数据组：\n', secret_group)

    # 解密部分
    # de_msg = decrypt(secret, key)
    # print('解密结果为:\n', de_msg)

    # 嵌入秘密信息
    packets = get_pool_packets(packet_num + 1)
    # print("packet_num=", len(packets))
    # print("secret_group=", len(secret_group))

    seq = 852353
    port = 12315
    pkt_list = []

    for i in range(packet_num):
        # packets[i].show2()
        seq += 20
        port += 1
        get_print(packets[i])
        
        temp = secret_insert(secret_group[i], seq, port, src, dst)
        pkt_list.append(temp)

    seq = 852350
    port = 12310
    get_print(packets[packet_num])
    pkt_end = secret_insert(hash_value, seq, port, src, dst, True)

    return packet_num, pkt_list, pkt_end

'''
